import backtrader as bt
import akshare as ak
import pandas as pd
from datetime import date, datetime, timedelta

def get_stocks_sector(symbol):
    print ('获取沪深300指数的成分股')
    hs300_df = ak.index_stock_cons(symbol=symbol) 
    hs300_stocks = hs300_df['品种代码'] 
    hs300_stocks=hs300_stocks.tolist()
    print('完成获取300股票代码')
    return hs300_stocks

# =========================
def last_day_of_month(year, month):
    """
    返回指定年份和月份的最后一天
    :param year: int, 年份
    :param month: int, 月份
    :return: date, 指定年份和月份的最后一天
    """
    # 获取下个月的第一天
    if month == 12:
        next_month_first_day = date(year + 1, 1, 1)
        print('下一个月第一天日期:',date(year + 1, 1, 1))
    else:
        next_month_first_day = date(year, month + 1, 1)
        print('下一个月第一天日期:',date(year, month + 1, 1))
    # 下个月的第一天减去一天就是本月的最后一天
    return next_month_first_day - timedelta(days=1)


# ====遍历300个股票,分别策略分析==========
def run_screener(stock_codes, start_date, end_date,num,head):
    print('加入数据源')
    red_days_stocks =pd.DataFrame(columns=['股票代码', '收红天数'])
    for stock_code in stock_codes:
        stock_data = ak.stock_zh_a_hist(symbol=stock_code, 
                                period="daily", 
                                start_date=start_date, 
                                end_date=end_date)
        red_days = stock_data[stock_data['收盘'] > stock_data['开盘']].shape[0]
        if red_days>=num:
            print(f"Red days for {stock_code}: {red_days}")
            new_row = pd.DataFrame({'股票代码': [stock_code], '收红天数': [red_days]})
            red_days_stocks = pd.concat([red_days_stocks, new_row], ignore_index=True)
            # 根据 '收红天数' 从高到低排序
            red_days_stocks = red_days_stocks.sort_values(by='收红天数', ascending=False)
            top_10_stocks = red_days_stocks.head(head)
    return top_10_stocks

if __name__ == '__main__':
    year = 2024
    month = 6
    symbol="000300"
    num=10  # 阳线根数
    head=10 # 符合条件前十名
    # today = datetime.now().date()
    today=last_day_of_month(year, month)
    today_str = today.strftime("%Y%m%d")
    today = datetime.strptime(today_str, "%Y%m%d")
    start_date=(today - timedelta(days=30)).strftime("%Y%m%d") 
    end_date=today.strftime("%Y%m%d")
    print(f"开始日期: {start_date}")
    # 获取板块股票
    stock_codes= get_stocks_sector(symbol)
    # 选出阳线大于几根的股票
    select_stock=run_screener(stock_codes=stock_codes, 
                                start_date=start_date, 
                                end_date=end_date, 
                                num=num,
                                head=head
                                )
    print(select_stock)